package main

import (
	"context"
	"fmt"
	"github.com/coreos/etcd/client"
	"github.com/coreos/etcd/pkg/transport"
	"math/rand"
	"strings"
	"sync"
	"time"
)

// Golang - [Iris] 服务注册与发现
// todo https://blog.csdn.net/wyansai/article/details/95157965
type ClientDis struct {
	client     client.KeysAPI
	serverList map[string]string
	lock       sync.Mutex
}

func NewClientDis(addr []string, prefixes []string) (cli *ClientDis, err error) {
	cli = &ClientDis{serverList: make(map[string]string)}

	tr, err := transport.NewTransport(transport.TLSInfo{}, 5*time.Second)
	if err != nil {
		return
	}

	conf := client.Config{Endpoints: addr, Transport: tr}

	var c client.Client
	if c, err = client.New(conf); err != nil {
		return
	}

	cli.client = client.NewKeysAPI(c)

	for _, p := range prefixes {
		resp, err := cli.client.Get(context.Background(), p, &client.GetOptions{Recursive: true})
		if err == nil && resp != nil && resp.Node.Dir && len(resp.Node.Nodes) > 0 {
			for _, n := range resp.Node.Nodes {
				if !n.Dir {
					cli.SetServiceList(n.Key, n.Value)
				}
			}
		}

		go func() {
			watcher := cli.client.Watcher(p, &client.WatcherOptions{Recursive: true, AfterIndex: 0})

			for {
				resp, err = watcher.Next(context.Background())

				if err != nil {
					fmt.Printf("%v", err)
					continue
				}

				if !resp.Node.Dir {
					switch resp.Action {
					case "set", "update":
						cli.SetServiceList(resp.Node.Key, resp.Node.Value)
					case "expire", "delete":
						cli.DelServiceList(resp.Node.Key)
					default:
						fmt.Println("watchMe!!!", "resp ->", resp)
					}
				}
			}

		}()
	}
	return
}

func (this *ClientDis) SetServiceList(key, val string) {
	this.lock.Lock()
	defer this.lock.Unlock()

	this.serverList[key] = string(val)
}

func (this *ClientDis) DelServiceList(key string) {
	this.lock.Lock()
	defer this.lock.Unlock()

	delete(this.serverList, key)
}

func (this *ClientDis) ServiceOne(prefix string) string {
	this.lock.Lock()
	defer this.lock.Unlock()

	addrArr := make([]string, 0, 16)
	for k, v := range this.serverList {
		if strings.HasPrefix(k, prefix) {
			addrArr = append(addrArr, v)
		}
	}

	if len(addrArr) == 0 {
		return ""
	}

	// 随机获取一个服务
	r := rand.New(rand.NewSource(time.Now().Unix()))
	return addrArr[r.Intn(len(addrArr))]
}

// 使用方法
func main() {
	// 服务发现
	cli, err := NewClientDis([]string{"127.0.0.1:2380"}, []string{"/node"})
	if err == nil {
		fmt.Print(cli.ServiceOne("/node"))
	} else {
		fmt.Printf("%v", err)
	}

	time.Sleep(10000 * time.Second)
}
